# -*- coding = utf-8 -*-
# @Time : 2022/1/15 21:19
# @Author : Chunyan Wei
# @File : predict.py
# @Software:PyCharm
import torch
import torchvision.transforms as transforms
from PIL import Image
from model import LeNet

def main():
    transform = transforms.Compose(
        [transforms.Resize((32,32)),
         transforms.ToTensor(),
         transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])
    classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    net = LeNet()
    net.load_state_dict(torch.load('Lenet.pth'))

    im = Image.open('1.png').convert('RGB')
    im = transform(im)
    im = torch.unsqueeze(im,dim=0)
    with torch.no_grad():
        outputs = net(im)
        predict = torch.max(outputs,dim=1)[1].data.numpy()
    print(classes[int(predict)])

if __name__ == '__main__':
    main()